import numpy as np
import pandas as pd
from python_ai.common.xcommon import sep
pd.set_option('display.max_columns', None)
from sklearn.feature_extraction.text import CountVectorizer

if '__main__' == __name__:
    X = ['我 爱 你',
         '我 恨 你 恨 你']

    count = CountVectorizer(token_pattern=r'\b\S+\b')
    tf = count.fit_transform(X)
    print(count.get_feature_names())
    print(tf.A)

    from sklearn.feature_extraction.text import TfidfVectorizer
    tf_idf_er = TfidfVectorizer(norm=None,  # ATTENTION
                                token_pattern=r'\b\S+\b')
    tf_idf = tf_idf_er.fit_transform(X)
    print(tf_idf_er.get_feature_names())
    print(tf_idf.A)

    sep('my tf-idf')


class TfidfMyImpl(object):

    def __init__(self, token_pattern=r'\b\S+\b'):
        self.token_pattern_ = token_pattern

    def fit_transform(self, list_of_sentences):
        tf_model = CountVectorizer(token_pattern=self.token_pattern_)
        tf = tf_model.fit_transform(list_of_sentences).A
        self.feature_names_ = tf_model.get_feature_names()
        tf = pd.DataFrame(tf, columns=self.feature_names_)
        # print(tf)
        idf = 1 + np.log((1 + len(list_of_sentences))/(1 + (tf != 0).sum(axis=0)))
        # print(idf)
        tf_idf = tf * idf
        # print(tf_idf)
        return tf_idf


    def get_feature_names(self):
        return self.feature_names_


if '__main__' == __name__:
    tf_idf_model = TfidfMyImpl()
    tf_idf = tf_idf_model.fit_transform(X)
    print(tf_idf)
